"""
XGBoost模型用于处理结构化特征
"""

import xgboost as xgb
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, classification_report, roc_auc_score
import joblib
import os

class XGBoostPredictor:
    """XGBoost预测器"""
    
    def __init__(self):
        self.hypertension_model = None
        self.diabetes_model = None
        self.is_trained = False
        
    def train(self, X, y_hypertension, y_diabetes):
        """
        训练XGBoost模型
        
        Args:
            X (DataFrame): 特征数据
            y_hypertension (array): 高血压标签
            y_diabetes (array): 糖尿病标签
        """
        try:
            # 分割训练和测试数据
            X_train, X_test, y_ht_train, y_ht_test = train_test_split(
                X, y_hypertension, test_size=0.2, random_state=42, stratify=y_hypertension
            )
            _, _, y_dm_train, y_dm_test = train_test_split(
                X, y_diabetes, test_size=0.2, random_state=42, stratify=y_diabetes
            )
            
            # 训练高血压预测模型
            self.hypertension_model = xgb.XGBClassifier(
                n_estimators=100,
                max_depth=6,
                learning_rate=0.1,
                subsample=0.8,
                colsample_bytree=0.8,
                random_state=42,
                eval_metric='logloss'
            )
            
            self.hypertension_model.fit(
                X_train, y_ht_train,
                eval_set=[(X_test, y_ht_test)],
                early_stopping_rounds=10,
                verbose=False
            )
            
            # 训练糖尿病预测模型
            self.diabetes_model = xgb.XGBClassifier(
                n_estimators=100,
                max_depth=6,
                learning_rate=0.1,
                subsample=0.8,
                colsample_bytree=0.8,
                random_state=42,
                eval_metric='logloss'
            )
            
            self.diabetes_model.fit(
                X_train, y_dm_train,
                eval_set=[(X_test, y_dm_test)],
                early_stopping_rounds=10,
                verbose=False
            )
            
            # 评估模型性能
            self._evaluate_models(X_test, y_ht_test, y_dm_test)
            
            self.is_trained = True
            print("XGBoost模型训练完成！")
            
        except Exception as e:
            print(f"XGBoost模型训练错误: {e}")
    
    def predict(self, X):
        """
        预测疾病风险
        
        Args:
            X (DataFrame or dict): 输入特征
            
        Returns:
            dict: 预测结果
        """
        if not self.is_trained:
            # 如果模型未训练，返回模拟预测结果
            return self._generate_mock_prediction(X)
        
        try:
            # 确保输入是DataFrame格式
            if isinstance(X, dict):
                X = pd.DataFrame([X])
            
            # 预测高血压风险
            ht_proba = self.hypertension_model.predict_proba(X)[0]
            ht_risk = ht_proba[1]  # 患病概率
            
            # 预测糖尿病风险
            dm_proba = self.diabetes_model.predict_proba(X)[0]
            dm_risk = dm_proba[1]  # 患病概率
            
            # 计算综合风险等级
            overall_risk = (ht_risk + dm_risk) / 2
            
            if overall_risk < 0.3:
                risk_level = 'low'
            elif overall_risk < 0.6:
                risk_level = 'medium'
            else:
                risk_level = 'high'
            
            return {
                'hypertension_risk': {
                    'probability': float(ht_risk),
                    'risk_level': 'high' if ht_risk > 0.5 else 'medium' if ht_risk > 0.3 else 'low'
                },
                'diabetes_risk': {
                    'probability': float(dm_risk),
                    'risk_level': 'high' if dm_risk > 0.5 else 'medium' if dm_risk > 0.3 else 'low'
                },
                'overall_risk': {
                    'probability': float(overall_risk),
                    'risk_level': risk_level
                },
                'prediction_confidence': float(max(ht_proba.max(), dm_proba.max()))
            }
            
        except Exception as e:
            print(f"XGBoost预测错误: {e}")
            return self._generate_mock_prediction(X)
    
    def _evaluate_models(self, X_test, y_ht_test, y_dm_test):
        """评估模型性能"""
        try:
            # 高血压模型评估
            ht_pred = self.hypertension_model.predict(X_test)
            ht_proba = self.hypertension_model.predict_proba(X_test)[:, 1]
            
            ht_accuracy = accuracy_score(y_ht_test, ht_pred)
            ht_auc = roc_auc_score(y_ht_test, ht_proba)
            
            print(f"高血压预测模型 - 准确率: {ht_accuracy:.3f}, AUC: {ht_auc:.3f}")
            
            # 糖尿病模型评估
            dm_pred = self.diabetes_model.predict(X_test)
            dm_proba = self.diabetes_model.predict_proba(X_test)[:, 1]
            
            dm_accuracy = accuracy_score(y_dm_test, dm_pred)
            dm_auc = roc_auc_score(y_dm_test, dm_proba)
            
            print(f"糖尿病预测模型 - 准确率: {dm_accuracy:.3f}, AUC: {dm_auc:.3f}")
            
        except Exception as e:
            print(f"模型评估错误: {e}")
    
    def _generate_mock_prediction(self, X):
        """生成模拟预测结果"""
        # 基于输入特征生成合理的模拟预测
        if isinstance(X, dict):
            age = X.get('age', 50)
            bmi = X.get('bmi', 23)
            systolic_bp = X.get('systolic_bp', 120)
            blood_glucose = X.get('blood_glucose', 100)
            smoking = X.get('smoking', 0)
            family_history = X.get('family_history', 0)
        else:
            # 如果是DataFrame，取第一行
            if len(X) > 0:
                age = X.iloc[0].get('age', 50)
                bmi = X.iloc[0].get('bmi', 23)
                systolic_bp = X.iloc[0].get('systolic_bp', 120)
                blood_glucose = X.iloc[0].get('blood_glucose', 100)
                smoking = X.iloc[0].get('smoking', 0)
                family_history = X.iloc[0].get('family_history', 0)
            else:
                return self._default_prediction()
        
        # 基于特征计算风险概率
        ht_risk = 0.1 + (age - 30) * 0.005 + (bmi - 20) * 0.01 + (systolic_bp - 120) * 0.002 + smoking * 0.1 + family_history * 0.15
        ht_risk = max(0.05, min(0.95, ht_risk))
        
        dm_risk = 0.05 + (age - 30) * 0.004 + (bmi - 20) * 0.015 + (blood_glucose - 100) * 0.003 + smoking * 0.08 + family_history * 0.12
        dm_risk = max(0.05, min(0.95, dm_risk))
        
        overall_risk = (ht_risk + dm_risk) / 2
        
        if overall_risk < 0.3:
            risk_level = 'low'
        elif overall_risk < 0.6:
            risk_level = 'medium'
        else:
            risk_level = 'high'
        
        return {
            'hypertension_risk': {
                'probability': float(ht_risk),
                'risk_level': 'high' if ht_risk > 0.5 else 'medium' if ht_risk > 0.3 else 'low'
            },
            'diabetes_risk': {
                'probability': float(dm_risk),
                'risk_level': 'high' if dm_risk > 0.5 else 'medium' if dm_risk > 0.3 else 'low'
            },
            'overall_risk': {
                'probability': float(overall_risk),
                'risk_level': risk_level
            },
            'prediction_confidence': 0.85
        }
    
    def _default_prediction(self):
        """默认预测结果"""
        return {
            'hypertension_risk': {
                'probability': 0.25,
                'risk_level': 'low'
            },
            'diabetes_risk': {
                'probability': 0.20,
                'risk_level': 'low'
            },
            'overall_risk': {
                'probability': 0.225,
                'risk_level': 'low'
            },
            'prediction_confidence': 0.80
        }
    
    def get_feature_importance(self):
        """获取特征重要性"""
        if not self.is_trained:
            return {}
        
        try:
            importance = {
                'hypertension': dict(zip(
                    self.hypertension_model.feature_names_in_,
                    self.hypertension_model.feature_importances_
                )),
                'diabetes': dict(zip(
                    self.diabetes_model.feature_names_in_,
                    self.diabetes_model.feature_importances_
                ))
            }
            return importance
        except Exception as e:
            print(f"获取特征重要性错误: {e}")
            return {}
    
    def save_models(self, filepath='models/xgboost_models.pkl'):
        """保存模型"""
        if self.is_trained:
            try:
                os.makedirs(os.path.dirname(filepath), exist_ok=True)
                joblib.dump({
                    'hypertension_model': self.hypertension_model,
                    'diabetes_model': self.diabetes_model
                }, filepath)
                print(f"模型已保存到: {filepath}")
            except Exception as e:
                print(f"保存模型错误: {e}")
    
    def load_models(self, filepath='models/xgboost_models.pkl'):
        """加载模型"""
        try:
            if os.path.exists(filepath):
                models = joblib.load(filepath)
                self.hypertension_model = models['hypertension_model']
                self.diabetes_model = models['diabetes_model']
                self.is_trained = True
                print(f"模型已从 {filepath} 加载")
            else:
                print(f"模型文件不存在: {filepath}")
        except Exception as e:
            print(f"加载模型错误: {e}")
